package com.small.xx.ai.spring.vector.controller;

import cn.hutool.core.io.FileUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 模型： ollama: mxbai-embed-large:latest
 */
@RestController
@RequestMapping("chat/vector/simple")
public class SimpleVectorController {
    private static final Logger logger = LoggerFactory.getLogger(SimpleVectorController.class);
    private final SimpleVectorStore simpleVectorStore;
    private final String SAVE_PATH = System.getProperty("user.dir") + "/doc/data/save.json";

    public SimpleVectorController(EmbeddingModel embeddingModel) {
        this.simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start add data");

        HashMap<String, Object> map = new HashMap<>();
        map.put("year", 2025);
        map.put("name", "XX");
        // 创建一个包含多个文档的列表，每个文档可能带有元数据
        List<Document> documents = List.of(new Document("世界很大， 救赎就潜伏在角落。"),
            new Document("你面向过去前行，又转身面向未来。", Map.of("year", 2024)),
            new Document("Spring AI太棒了！Spring AI太棒了！Spring AI太棒了！Spring AI太棒了！Spring AI太棒了！", map),
            new Document("其他", "测试内容", map));
        // 将文档列表添加到简单的向量存储中，用于后续的检索或处理
        simpleVectorStore.add(documents);
    }

    @GetMapping("/delete")
    public void delete() {
        logger.info("start delete data");
        simpleVectorStore.delete(List.of("1"));
    }

    @GetMapping("/save")
    public void save() {
        logger.info("start save data: {}", SAVE_PATH);
        File file = new File(SAVE_PATH);
        if (file.exists()) {
            file.delete();
        }
        FileUtil.mkdir(file.getParent());
        simpleVectorStore.save(file);
    }

    @GetMapping("/load")
    public void load() {
        logger.info("start load data: {}", SAVE_PATH);
        File file = new File(SAVE_PATH);
        simpleVectorStore.load(file);
    }

    @GetMapping("/search")
    public List<Document> search() {
        logger.info("start search data");
        return simpleVectorStore.similaritySearch(SearchRequest.builder().query("Spring").topK(2).build());
    }

    @GetMapping("/search-filter")
    public List<Document> searchFilter() {
        logger.info("start search  filter data");
        FilterExpressionBuilder b = new FilterExpressionBuilder();
//        Filter.Expression expression = b.and(b.in("year", 2025, 2024), b.eq("name", "XX")).build();
        Filter.Expression expression = b.or(b.in("year", 2024), b.eq("name", "XX")).build();

        return simpleVectorStore.similaritySearch(
            SearchRequest.builder().query("Spring").topK(2).filterExpression(expression).build());
    }
}
